import os
import re
import pandas as pd
from consts import (
  FEATURE_KEYWORD,
  FEATURE_GENRE,
  FEATURE_INFO,
  FEATURE_ALL,
)
from utils import (
  log_message,
  is_numeric,
  round_to_nearest_half
)

def merge_dir_path(feature_type, is_input=True):
  if is_input:
    path = f'../data/result/basic_large_model/origin_data/{feature_type}'
  else:
    path = f'../data/result/basic_large_model/final_data/{feature_type}'
  return path

# 处理评分数据
def process_result_score(feature_type, filename):
  input_directory_path = merge_dir_path(feature_type, is_input=True)
  output_directory_path = merge_dir_path(feature_type, is_input=False)
  file_path = os.path.join(input_directory_path, filename)
  output_path = os.path.join(output_directory_path, filename)
  str_arr = filename.split('.')
  if len(str_arr) > 2:
    filename = ''.join(str_arr[:len(str_arr) - 1])
  else:
    filename = str_arr[0]
  # 读取CSV文件
  df = pd.read_csv(file_path)
  data = df.to_dict(orient='records')
  """
  这里要处理两种异常的数据
  1.有些模型太弱了，给了提示词，给出的回复还是一堆文字，但最终也是给了得分，因此我们需要通过正则表达式获取数据
  例如：Based on the user's past rating history and the keywords of the candidate movie "I.Q.", I predict the user rating to be 3.5.
  匹配3.5
  2.有些数据给出的评分分数很奇怪，比如3.8、2.4等，我们需要的是[0, 0.5, 1.0, ..., 4.5, 5.0]
  需要将这些数据转换成我们需要的评分数据
  """
  for item in data:
    if not is_numeric(item['predict']):
      # 正则匹配
      pattern = r'(\d\.\d)'
      matchs = re.findall(pattern, item['predict'])
      if len(matchs) == 0:
        score = 2.5
      else:
        score = round_to_nearest_half(float(matchs[0]))
      item['predict'] = str(score)
  result = pd.DataFrame(data)
  result.to_csv(output_path, index=False)

def process_dir_file(feature_type):
  input_path = merge_dir_path(feature_type, is_input=True)
  for filename in os.listdir(input_path):
    # 检查文件是否为CSV文件
    if filename.endswith('.csv'):
      log_message(f"文件名称{filename}")
      process_result_score(feature_type, filename)
      log_message("处理完毕")
      # 执行一些操作，例如打印前五行

def main():
  process_dir_file(FEATURE_KEYWORD)
  process_dir_file(FEATURE_GENRE)
  process_dir_file(FEATURE_INFO)
  process_dir_file(FEATURE_ALL)
main()
